// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/common/process/trace_context.h>

#include <pollux/common/process/trace_history.h>

#include <sstream>

namespace kumo::pollux::process {
    namespace {
        // We use thread local instead lock here since the critical path is on write
        // side.
        auto registry = std::make_shared<TraceContext::Registry>();
        thread_local auto threadLocalTraceData =
                std::make_shared<TraceContext::Registry::Reference>(registry);
    } // namespace

    TraceContext::TraceContext(std::string label, bool isTemporary)
        : label_(std::move(label)),
          enterTime_(std::chrono::steady_clock::now()),
          isTemporary_(isTemporary),
          traceData_(threadLocalTraceData) {
        TraceHistory::push([&](auto &entry) {
            entry.time = enterTime_;
            entry.file = __FILE__;
            entry.line = __LINE__;
            snprintf(entry.label, entry.kLabelCapacity, "%s", label_.c_str());
        });
        traceData_->withValue([&](auto &counts) {
            auto &data = counts[label_];
            ++data.numThreads;
            if (data.numThreads == 1) {
                data.startTime = enterTime_;
            }
            ++data.numEnters;
        });
    }

    TraceContext::~TraceContext() {
        traceData_->withValue([&](auto &counts) {
            auto it = counts.find(label_);
            auto &data = it->second;
            if (--data.numThreads == 0 && isTemporary_) {
                counts.erase(it);
                return;
            }
            auto ms = std::chrono::duration_cast<std::chrono::milliseconds>(
                        std::chrono::steady_clock::now() - enterTime_)
                    .count();
            data.totalMs += ms;
            data.maxMs = std::max<uint64_t>(data.maxMs, ms);
        });
    }

    // static
    std::string TraceContext::statusLine() {
        std::stringstream out;
        auto now = std::chrono::steady_clock::now();
        auto counts = status();
        for (auto &[label, data]: counts) {
            if (data.numThreads > 0) {
                auto continued = std::chrono::duration_cast<std::chrono::milliseconds>(
                            now - data.startTime)
                        .count();
                out << label << ": numThreads=" << data.numThreads
                        << " numEnters=" << data.numEnters
                        << " avgMs=" << (data.totalMs / data.numEnters)
                        << " maxMs=" << data.maxMs << " continued=" << continued << std::endl;
            }
        }
        return out.str();
    }

    // static
    melon::F14FastMap<std::string, TraceData> TraceContext::status() {
        melon::F14FastMap<std::string, TraceData> total;
        registry->forAllValues([&](auto &counts) {
            for (auto &[k, v]: counts) {
                auto &sofar = total[k];
                if (sofar.numEnters == 0) {
                    sofar.startTime = v.startTime;
                } else if (v.numEnters > 0) {
                    sofar.startTime = std::min(sofar.startTime, v.startTime);
                }
                sofar.numThreads += v.numThreads;
                sofar.numEnters += v.numEnters;
                sofar.totalMs += v.totalMs;
                sofar.maxMs = std::max(sofar.maxMs, v.maxMs);
            }
        });
        return total;
    }
} // namespace kumo::pollux::process
